# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import typing

from pydantic import Field
from pydantic import model_validator

from aiq.data_models.component_ref import LLMRef
from aiq.data_models.its_strategy import ITSStrategyBaseConfig


class LLMBasedPlanScoringConfig(ITSStrategyBaseConfig, name="llm_based_plan_scoring"):
    """
    Configuration for LLMBasedScoring.
    """
    scoring_llm: LLMRef | typing.Any | None = Field(
        default=None,
        description="The LLM to use for scoring the plans. This can be a callable or an instance of an LLM client.")

    scoring_template: str = Field(
        default=("You are an expert reasoning model tasked with scoring the following execution plan based on its"
                 "quality and relevance to the provided input to an agent system.\n\n"
                 "The agent system's role is:\n{context}\n\n"
                 "It has been tasked with achieving the following goal: \n{original_prompt}\n\n"
                 "The following plan has been generated to achieve this goal:\n\n{plan}\n\n"
                 "Score the plan on a scale from 1 to 10, where 10 is the best. "
                 "Return the final score as a floating point number preceded by `FINAL SCORE:` without any "
                 "other text before or after it\n"),
        description="The template to use for scoring the plans.")

    @model_validator(mode="before")
    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
        """
        Ensure that the scoring_llm is provided when using LLMBasedScoring.
        """
        if values.get('scoring_llm') is None:
            raise ValueError('scoring_llm must be provided when scorer_type is set to LLM_BASED_SCORING.')

        return values


class LLMBasedAgentScoringConfig(ITSStrategyBaseConfig, name="llm_based_agent_scoring"):
    """
    Configuration for LLMBasedScoring.
    """
    scoring_llm: LLMRef | typing.Any | None = Field(
        default=None,
        description="The LLM to use for scoring the plans. This can be a callable or an instance of an LLM client.")

    scoring_template: str = Field(
        description="Prompt template to use for scoring the function output",
        default=("You are an expert reasoning model tasked with scoring the following "
                 "result of an agent system based on its input and objective. Judge"
                 " the quality and relevance of the answer to score it.\n\n"
                 "The agent system's objective is:\n{objective}\n\n"
                 "It has been tasked with achieving the following goal: \n{input}\n\n"
                 "The following output has been generated by the agent:\n\n{output}\n\n"
                 "Score the result on a scale from 1 to 10, where 10 is the best. "
                 "Return the final score as a floating point number preceded by `FINAL SCORE:` without any "
                 "other text before or after it\n"),
    )

    @model_validator(mode="before")
    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
        """
        Ensure that the scoring_llm is provided when using LLMBasedScoring.
        """
        if values.get('scoring_llm') is None:
            raise ValueError('scoring_llm must be provided when scorer_type is set to LLM_BASED_SCORING.')

        return values


class MotivationAwareScoringConfig(ITSStrategyBaseConfig, name="motivation_aware_scoring"):
    """
    Configuration for a scoring strategy that considers both the original input (task)
    and the motivation (from metadata) along with the current output.
    """

    scoring_llm: LLMRef | None = Field(
        default=None, description="The LLM used to evaluate how well the output addresses the task plus motivation.")

    scoring_template: str = Field(
        default=("You are an expert at assessing the quality of an output in relation to its task and motivation.\n"
                 "Task: {task}\n"
                 "Motivation: {motivation}\n"
                 "Output: {output}\n"
                 "On a scale from 1 to 10 (10 being the best), how well does this output fulfill "
                 "the original task in the context "
                 "of the provided motivation? Note that the task might answer one part of a bigger question "
                 "which should count as a satisfactory response and should not receive a lower score.\n"
                 "Return the final score as a floating point number preceded by 'FINAL SCORE:'."),
        description="The prompt template used to evaluate and score the output.")

    @model_validator(mode="before")
    def validate_scoring_llm(cls, values):
        if values.get('scoring_llm') is None:
            raise ValueError("A scoring_llm must be provided for motivation_aware_scoring.")
        return values
